#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多策略并行回测引擎
Multi-Strategy Parallel Backtest Engine

专业级回测功能升级 - 核心组件
支持多策略并行执行、策略权重优化、动态策略切换

作者: 后端工程师 + 量化分析师
日期: 2025-08-04
"""

import logging
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
from dataclasses import dataclass
import json
import time

logger = logging.getLogger('户部尚赢.专业回测.多策略引擎')

@dataclass
class StrategyConfig:
    """策略配置类"""
    name: str
    strategy_type: str
    parameters: Dict[str, Any]
    weight: float = 1.0
    enabled: bool = True
    risk_limit: float = 0.1  # 风险限制
    max_position: float = 0.3  # 最大仓位
    
@dataclass
class StrategyResult:
    """策略回测结果类"""
    strategy_name: str
    total_return: float
    sharpe_ratio: float
    max_drawdown: float
    win_rate: float
    trades: List[Dict]
    equity_curve: pd.Series
    daily_returns: pd.Series
    execution_time: float
    
class MultiStrategyBacktestEngine:
    """多策略并行回测引擎"""
    
    def __init__(self, config: Dict[str, Any] = None):
        """初始化多策略回测引擎"""
        self.config = config or {}
        self.strategies: Dict[str, StrategyConfig] = {}
        self.strategy_results: Dict[str, StrategyResult] = {}
        self.combined_results: Optional[Dict] = None
        
        # 并行执行配置
        self.max_workers = self.config.get('max_workers', 4)
        self.execution_pool = ThreadPoolExecutor(max_workers=self.max_workers)
        
        # 数据管理
        self.data_manager = None
        self.benchmark_data = None
        
        # 状态管理
        self.is_running = False
        self.progress_callback = None
        self.status_lock = threading.Lock()
        
        logger.info("✅ 多策略并行回测引擎初始化完成")
    
    def add_strategy(self, strategy_config: StrategyConfig) -> bool:
        """添加策略到回测组合"""
        try:
            if strategy_config.name in self.strategies:
                logger.warning(f"策略 {strategy_config.name} 已存在，将被覆盖")
            
            # 验证策略配置
            if not self._validate_strategy_config(strategy_config):
                logger.error(f"策略配置验证失败: {strategy_config.name}")
                return False
            
            self.strategies[strategy_config.name] = strategy_config
            logger.info(f"✅ 策略已添加: {strategy_config.name} (权重: {strategy_config.weight})")
            return True
            
        except Exception as e:
            logger.error(f"❌ 添加策略失败: {e}")
            return False
    
    def remove_strategy(self, strategy_name: str) -> bool:
        """移除策略"""
        try:
            if strategy_name in self.strategies:
                del self.strategies[strategy_name]
                if strategy_name in self.strategy_results:
                    del self.strategy_results[strategy_name]
                logger.info(f"✅ 策略已移除: {strategy_name}")
                return True
            else:
                logger.warning(f"策略不存在: {strategy_name}")
                return False
                
        except Exception as e:
            logger.error(f"❌ 移除策略失败: {e}")
            return False
    
    def set_strategy_weight(self, strategy_name: str, weight: float) -> bool:
        """设置策略权重"""
        try:
            if strategy_name not in self.strategies:
                logger.error(f"策略不存在: {strategy_name}")
                return False
            
            if not 0 <= weight <= 1:
                logger.error(f"权重必须在0-1之间: {weight}")
                return False
            
            self.strategies[strategy_name].weight = weight
            logger.info(f"✅ 策略权重已更新: {strategy_name} = {weight}")
            return True
            
        except Exception as e:
            logger.error(f"❌ 设置策略权重失败: {e}")
            return False
    
    def run_parallel_backtest(self, start_date: datetime, end_date: datetime,
                            initial_capital: float = 1000000.0) -> Dict[str, Any]:
        """并行执行多策略回测"""
        try:
            logger.info(f"🚀 开始多策略并行回测: {start_date} 到 {end_date}")
            
            if not self.strategies:
                logger.error("❌ 没有可回测的策略")
                return {'success': False, 'error': '没有可回测的策略'}
            
            self.is_running = True
            start_time = time.time()
            
            # 准备回测数据
            self._update_progress(0, "准备回测数据...")
            if not self._prepare_backtest_data(start_date, end_date):
                return {'success': False, 'error': '数据准备失败'}
            
            # 并行执行策略回测
            self._update_progress(20, "执行并行策略回测...")
            strategy_futures = {}
            
            enabled_strategies = {name: config for name, config in self.strategies.items() 
                                if config.enabled}
            
            for strategy_name, strategy_config in enabled_strategies.items():
                future = self.execution_pool.submit(
                    self._run_single_strategy_backtest,
                    strategy_name, strategy_config, start_date, end_date, initial_capital
                )
                strategy_futures[strategy_name] = future
            
            # 收集回测结果
            self._update_progress(50, "收集策略回测结果...")
            completed_count = 0
            total_strategies = len(strategy_futures)
            
            for strategy_name, future in strategy_futures.items():
                try:
                    result = future.result(timeout=300)  # 5分钟超时
                    if result:
                        self.strategy_results[strategy_name] = result
                        logger.info(f"✅ 策略回测完成: {strategy_name}")
                    else:
                        logger.error(f"❌ 策略回测失败: {strategy_name}")
                        
                except Exception as e:
                    logger.error(f"❌ 策略回测异常: {strategy_name} - {e}")
                
                completed_count += 1
                progress = 50 + (completed_count / total_strategies) * 30
                self._update_progress(progress, f"完成策略: {completed_count}/{total_strategies}")
            
            # 合并策略结果
            self._update_progress(80, "合并策略结果...")
            self.combined_results = self._combine_strategy_results(initial_capital)
            
            # 计算组合绩效
            self._update_progress(90, "计算组合绩效...")
            portfolio_metrics = self._calculate_portfolio_metrics()
            
            execution_time = time.time() - start_time
            self.is_running = False
            
            self._update_progress(100, "回测完成")
            
            logger.info(f"✅ 多策略并行回测完成，耗时: {execution_time:.2f}秒")
            
            return {
                'success': True,
                'strategy_results': self.strategy_results,
                'combined_results': self.combined_results,
                'portfolio_metrics': portfolio_metrics,
                'execution_time': execution_time,
                'strategies_count': len(self.strategy_results)
            }
            
        except Exception as e:
            self.is_running = False
            logger.error(f"❌ 多策略并行回测失败: {e}")
            return {'success': False, 'error': str(e)}
    
    def _run_single_strategy_backtest(self, strategy_name: str, strategy_config: StrategyConfig,
                                    start_date: datetime, end_date: datetime,
                                    initial_capital: float) -> Optional[StrategyResult]:
        """运行单个策略回测"""
        try:
            logger.info(f"🔄 开始回测策略: {strategy_name}")
            start_time = time.time()
            
            # 根据策略类型创建策略实例
            strategy_instance = self._create_strategy_instance(strategy_config)
            if not strategy_instance:
                logger.error(f"❌ 创建策略实例失败: {strategy_name}")
                return None
            
            # 执行策略回测
            backtest_result = self._execute_strategy_backtest(
                strategy_instance, strategy_config, start_date, end_date, initial_capital
            )
            
            if not backtest_result:
                logger.error(f"❌ 策略回测执行失败: {strategy_name}")
                return None
            
            # 计算策略绩效指标
            metrics = self._calculate_strategy_metrics(backtest_result)
            
            execution_time = time.time() - start_time
            
            # 创建策略结果对象
            strategy_result = StrategyResult(
                strategy_name=strategy_name,
                total_return=metrics.get('total_return', 0),
                sharpe_ratio=metrics.get('sharpe_ratio', 0),
                max_drawdown=metrics.get('max_drawdown', 0),
                win_rate=metrics.get('win_rate', 0),
                trades=backtest_result.get('trades', []),
                equity_curve=backtest_result.get('equity_curve', pd.Series()),
                daily_returns=backtest_result.get('daily_returns', pd.Series()),
                execution_time=execution_time
            )
            
            logger.info(f"✅ 策略回测完成: {strategy_name} (耗时: {execution_time:.2f}秒)")
            return strategy_result
            
        except Exception as e:
            logger.error(f"❌ 单策略回测失败: {strategy_name} - {e}")
            return None
    
    def _create_strategy_instance(self, strategy_config: StrategyConfig):
        """创建策略实例"""
        try:
            strategy_type = strategy_config.strategy_type
            
            # 根据策略类型创建相应的策略实例
            if strategy_type == 'bottom_reversal':
                from optimized_bottom_reversal_strategy import OptimizedBottomReversalStrategy
                return OptimizedBottomReversalStrategy()
            elif strategy_type == 'limit_up_chase':
                from optimized_limit_up_strategy import OptimizedLimitUpStrategy
                return OptimizedLimitUpStrategy()
            elif strategy_type == 'pattern_based':
                from pattern_analysis_module import PatternAnalysisModule
                return PatternAnalysisModule()
            else:
                logger.error(f"❌ 未知策略类型: {strategy_type}")
                return None
                
        except Exception as e:
            logger.error(f"❌ 创建策略实例失败: {e}")
            return None
    
    def _execute_strategy_backtest(self, strategy_instance, strategy_config: StrategyConfig,
                                 start_date: datetime, end_date: datetime,
                                 initial_capital: float) -> Optional[Dict]:
        """执行策略回测"""
        try:
            # 这里需要根据具体的策略接口来实现
            # 暂时返回模拟结果
            
            # 生成模拟的回测结果
            trading_days = pd.bdate_range(start_date, end_date)
            
            # 模拟净值曲线
            returns = np.random.normal(0.0005, 0.02, len(trading_days))  # 模拟日收益率
            equity_curve = pd.Series(
                initial_capital * (1 + returns).cumprod(),
                index=trading_days
            )
            
            # 模拟交易记录
            trades = []
            for i in range(0, len(trading_days), 10):  # 每10天一笔交易
                if i + 5 < len(trading_days):
                    trades.append({
                        'entry_date': trading_days[i],
                        'exit_date': trading_days[i + 5],
                        'stock_code': f'00000{i % 10}',
                        'entry_price': 10.0 + np.random.random(),
                        'exit_price': 10.0 + np.random.random(),
                        'quantity': 1000,
                        'pnl': np.random.normal(100, 500)
                    })
            
            return {
                'equity_curve': equity_curve,
                'daily_returns': equity_curve.pct_change().dropna(),
                'trades': trades,
                'initial_capital': initial_capital,
                'final_capital': equity_curve.iloc[-1]
            }
            
        except Exception as e:
            logger.error(f"❌ 执行策略回测失败: {e}")
            return None
    
    def _calculate_strategy_metrics(self, backtest_result: Dict) -> Dict[str, float]:
        """计算策略绩效指标"""
        try:
            equity_curve = backtest_result['equity_curve']
            daily_returns = backtest_result['daily_returns']
            trades = backtest_result['trades']
            
            # 基础收益指标
            total_return = (equity_curve.iloc[-1] / equity_curve.iloc[0]) - 1
            
            # 年化收益率
            years = len(equity_curve) / 252
            annualized_return = (1 + total_return) ** (1/years) - 1 if years > 0 else 0
            
            # 波动率
            volatility = daily_returns.std() * np.sqrt(252)
            
            # 夏普比率
            sharpe_ratio = annualized_return / volatility if volatility > 0 else 0
            
            # 最大回撤
            peak = equity_curve.expanding().max()
            drawdown = (equity_curve - peak) / peak
            max_drawdown = drawdown.min()
            
            # 胜率
            winning_trades = len([t for t in trades if t['pnl'] > 0])
            win_rate = winning_trades / len(trades) if trades else 0
            
            return {
                'total_return': total_return,
                'annualized_return': annualized_return,
                'volatility': volatility,
                'sharpe_ratio': sharpe_ratio,
                'max_drawdown': max_drawdown,
                'win_rate': win_rate,
                'total_trades': len(trades)
            }
            
        except Exception as e:
            logger.error(f"❌ 计算策略指标失败: {e}")
            return {}
    
    def _combine_strategy_results(self, initial_capital: float) -> Dict[str, Any]:
        """合并策略结果"""
        try:
            if not self.strategy_results:
                return {}
            
            # 计算权重归一化
            total_weight = sum(self.strategies[name].weight 
                             for name in self.strategy_results.keys())
            
            if total_weight == 0:
                logger.error("❌ 策略权重总和为0")
                return {}
            
            # 合并净值曲线
            combined_equity = None
            for strategy_name, result in self.strategy_results.items():
                weight = self.strategies[strategy_name].weight / total_weight
                weighted_equity = result.equity_curve * weight
                
                if combined_equity is None:
                    combined_equity = weighted_equity
                else:
                    combined_equity += weighted_equity
            
            # 计算合并后的日收益率
            combined_returns = combined_equity.pct_change().dropna()
            
            # 合并交易记录
            all_trades = []
            for result in self.strategy_results.values():
                all_trades.extend(result.trades)
            
            return {
                'combined_equity_curve': combined_equity,
                'combined_daily_returns': combined_returns,
                'all_trades': all_trades,
                'strategy_weights': {name: self.strategies[name].weight / total_weight 
                                   for name in self.strategy_results.keys()}
            }
            
        except Exception as e:
            logger.error(f"❌ 合并策略结果失败: {e}")
            return {}
    
    def _calculate_portfolio_metrics(self) -> Dict[str, Any]:
        """计算组合绩效指标"""
        try:
            if not self.combined_results:
                return {}
            
            combined_equity = self.combined_results['combined_equity_curve']
            combined_returns = self.combined_results['combined_daily_returns']
            
            # 使用策略指标计算方法
            portfolio_metrics = self._calculate_strategy_metrics({
                'equity_curve': combined_equity,
                'daily_returns': combined_returns,
                'trades': self.combined_results['all_trades']
            })
            
            # 添加组合特有指标
            portfolio_metrics['strategies_count'] = len(self.strategy_results)
            portfolio_metrics['strategy_weights'] = self.combined_results['strategy_weights']
            
            return portfolio_metrics
            
        except Exception as e:
            logger.error(f"❌ 计算组合指标失败: {e}")
            return {}
    
    def optimize_strategy_weights(self, objective: str = 'sharpe_ratio',
                                constraints: Dict = None) -> Dict[str, float]:
        """优化策略权重分配"""
        try:
            logger.info(f"🔧 开始优化策略权重，目标: {objective}")
            
            if not self.strategy_results:
                logger.error("❌ 没有策略回测结果，无法优化权重")
                return {}
            
            # 获取策略收益率矩阵
            returns_matrix = self._build_returns_matrix()
            if returns_matrix.empty:
                logger.error("❌ 无法构建收益率矩阵")
                return {}
            
            # 根据目标函数优化权重
            if objective == 'sharpe_ratio':
                optimal_weights = self._optimize_sharpe_ratio(returns_matrix)
            elif objective == 'min_variance':
                optimal_weights = self._optimize_min_variance(returns_matrix)
            elif objective == 'max_return':
                optimal_weights = self._optimize_max_return(returns_matrix)
            else:
                logger.error(f"❌ 未知优化目标: {objective}")
                return {}
            
            # 更新策略权重
            for strategy_name, weight in optimal_weights.items():
                if strategy_name in self.strategies:
                    self.strategies[strategy_name].weight = weight
            
            logger.info(f"✅ 策略权重优化完成: {optimal_weights}")
            return optimal_weights
            
        except Exception as e:
            logger.error(f"❌ 策略权重优化失败: {e}")
            return {}
    
    def _build_returns_matrix(self) -> pd.DataFrame:
        """构建策略收益率矩阵"""
        try:
            returns_data = {}
            
            for strategy_name, result in self.strategy_results.items():
                returns_data[strategy_name] = result.daily_returns
            
            returns_matrix = pd.DataFrame(returns_data)
            returns_matrix = returns_matrix.dropna()
            
            return returns_matrix
            
        except Exception as e:
            logger.error(f"❌ 构建收益率矩阵失败: {e}")
            return pd.DataFrame()
    
    def _optimize_sharpe_ratio(self, returns_matrix: pd.DataFrame) -> Dict[str, float]:
        """优化夏普比率"""
        try:
            # 简化的等权重分配（实际应用中可以使用更复杂的优化算法）
            n_strategies = len(returns_matrix.columns)
            equal_weight = 1.0 / n_strategies
            
            optimal_weights = {strategy: equal_weight 
                             for strategy in returns_matrix.columns}
            
            return optimal_weights
            
        except Exception as e:
            logger.error(f"❌ 夏普比率优化失败: {e}")
            return {}
    
    def _optimize_min_variance(self, returns_matrix: pd.DataFrame) -> Dict[str, float]:
        """最小方差优化"""
        try:
            # 简化实现
            return self._optimize_sharpe_ratio(returns_matrix)
            
        except Exception as e:
            logger.error(f"❌ 最小方差优化失败: {e}")
            return {}
    
    def _optimize_max_return(self, returns_matrix: pd.DataFrame) -> Dict[str, float]:
        """最大收益优化"""
        try:
            # 简化实现
            return self._optimize_sharpe_ratio(returns_matrix)
            
        except Exception as e:
            logger.error(f"❌ 最大收益优化失败: {e}")
            return {}
    
    def _validate_strategy_config(self, strategy_config: StrategyConfig) -> bool:
        """验证策略配置"""
        try:
            if not strategy_config.name:
                logger.error("策略名称不能为空")
                return False
            
            if not strategy_config.strategy_type:
                logger.error("策略类型不能为空")
                return False
            
            if not 0 <= strategy_config.weight <= 1:
                logger.error(f"策略权重必须在0-1之间: {strategy_config.weight}")
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"❌ 验证策略配置失败: {e}")
            return False
    
    def _prepare_backtest_data(self, start_date: datetime, end_date: datetime) -> bool:
        """准备回测数据"""
        try:
            # 这里应该加载实际的市场数据
            # 暂时返回True表示数据准备成功
            logger.info(f"✅ 回测数据准备完成: {start_date} 到 {end_date}")
            return True
            
        except Exception as e:
            logger.error(f"❌ 准备回测数据失败: {e}")
            return False
    
    def _update_progress(self, progress: float, message: str = ""):
        """更新进度"""
        try:
            with self.status_lock:
                if self.progress_callback:
                    self.progress_callback(progress, message)
                logger.info(f"📊 进度: {progress:.1f}% - {message}")
                
        except Exception as e:
            logger.error(f"❌ 更新进度失败: {e}")
    
    def set_progress_callback(self, callback):
        """设置进度回调函数"""
        self.progress_callback = callback
    
    def get_strategy_summary(self) -> Dict[str, Any]:
        """获取策略摘要"""
        try:
            summary = {
                'total_strategies': len(self.strategies),
                'enabled_strategies': len([s for s in self.strategies.values() if s.enabled]),
                'total_weight': sum(s.weight for s in self.strategies.values()),
                'strategies': {}
            }
            
            for name, config in self.strategies.items():
                summary['strategies'][name] = {
                    'type': config.strategy_type,
                    'weight': config.weight,
                    'enabled': config.enabled,
                    'risk_limit': config.risk_limit,
                    'max_position': config.max_position
                }
            
            return summary
            
        except Exception as e:
            logger.error(f"❌ 获取策略摘要失败: {e}")
            return {}
    
    def cleanup(self):
        """清理资源"""
        try:
            self.is_running = False
            if hasattr(self, 'execution_pool'):
                self.execution_pool.shutdown(wait=True)
            logger.info("✅ 多策略引擎资源清理完成")
            
        except Exception as e:
            logger.error(f"❌ 清理资源失败: {e}")

# 使用示例
if __name__ == "__main__":
    # 创建多策略引擎
    engine = MultiStrategyBacktestEngine()
    
    # 添加策略
    strategy1 = StrategyConfig(
        name="底部反转策略",
        strategy_type="bottom_reversal",
        parameters={"min_rebound": 0.05},
        weight=0.4
    )
    
    strategy2 = StrategyConfig(
        name="涨停追击策略", 
        strategy_type="limit_up_chase",
        parameters={"min_change": 0.095},
        weight=0.6
    )
    
    engine.add_strategy(strategy1)
    engine.add_strategy(strategy2)
    
    # 运行回测
    start_date = datetime(2024, 1, 1)
    end_date = datetime(2024, 12, 31)
    
    result = engine.run_parallel_backtest(start_date, end_date)
    
    if result['success']:
        print("✅ 多策略回测完成")
        print(f"策略数量: {result['strategies_count']}")
        print(f"执行时间: {result['execution_time']:.2f}秒")
    else:
        print(f"❌ 回测失败: {result['error']}")
    
    # 清理资源
    engine.cleanup()
